Skip to content

Add new BOUNDS_CHECK GatherScatterMode #33446

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add new BOUNDS_CHECK GatherScatterMode #33446

wants to merge 1 commit into from

Conversation

@jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Nov 20, 2025

No description provided.

@jakevdp jakevdp self-assigned this Nov 20, 2025
@gemini-code-assist
Copy link
Contributor

gemini-code-assist bot commented Nov 20, 2025

Summary of Changes

Hello @jakevdp, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new BOUNDS_CHECK mode for JAX's gather and scatter operations, enhancing the robustness and debuggability of array indexing. This mode explicitly raises an IndexError for out-of-bounds static indices, providing immediate feedback to users about invalid access patterns. This prevents silent data manipulation that could occur with other modes like FILL_OR_DROP when indices are outside the array boundaries. The changes involve updating the GatherScatterMode enum, integrating this new behavior into the core indexing logic, and ensuring proper test coverage and compatibility with existing performance models.

Highlights

  • New GatherScatterMode: BOUNDS_CHECK: Introduced a new GatherScatterMode called BOUNDS_CHECK. This mode is designed to raise an IndexError when out-of-bounds indices are detected, particularly for static indices or when used with jax.experimental.checkify. For dynamic indices, its behavior defaults to FILL_OR_DROP.
  • Integration into Gather and Scatter Operations: The new BOUNDS_CHECK mode has been integrated into the core gather and scatter functions across jax/_src/lax/slicing.py, jax/_src/numpy/indexing.py, and jax/_src/ops/scatter.py to enable its error-raising behavior.
  • Static Index Bounds Checking: Logic has been added to jax/_src/numpy/indexing.py to perform explicit static bounds checking. When BOUNDS_CHECK mode is active, index_to_gather will now raise an IndexError if a static index falls outside the array's dimensions.
  • Updated Roofline Model: The roofline performance model in jax/experimental/roofline/rooflines.py and its associated tests have been updated to correctly account for the BOUNDS_CHECK mode, treating its FLOPs calculation similarly to FILL_OR_DROP.
  • Comprehensive Test Coverage: New parameterized tests have been added in tests/lax_numpy_indexing_test.py to thoroughly validate the BOUNDS_CHECK mode, ensuring that IndexError is correctly raised for various out-of-bounds static indexing scenarios in both gather and scatter operations.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist[bot]
gemini-code-assist bot reviewed Nov 20, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a new BOUNDS_CHECK mode for GatherScatterMode. This is a valuable addition that will help catch out-of-bounds indexing errors at compile time for static indices. The implementation is well-integrated across the relevant parts of JAX, including lax.slicing, numpy.indexing, and ops.scatter. The fallback behavior to FILL_OR_DROP when static checking is not possible is sensible. The new functionality is also thoroughly tested.

I've made a couple of minor suggestions to use tuples instead of lists for in checks for consistency across the codebase. Overall, this is a great change.

@jakevdp jakevdp requested a review from mattjj November 20, 2025 20:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mattjj mattjj Awaiting requested review from mattjj

1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

@jakevdp jakevdp

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@jakevdp